#========================================
# author: changlong.zang
#   mail: zclongpop@163.com
#   date: Fri, 26 Sep 2014 15:50:26
#========================================
import os, re, json, scriptTool
import maya.cmds as mc
import maya.mel as mel
#--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
def selectDictionary():
    '''
    select dir path...
    '''
    result = ''
    path = mc.fileDialog2(fm=3, okc='Select')
    if path:
        result = path[0]
    return result


def selectFile(startPath='', fullPath=False):
    '''
    Select file path from given path...
    '''
    if os.path.isfile(startPath):
        dir_path = os.path.dirname(startPath)
    elif os.path.isdir(startPath):
        dir_path = startPath
    else:
        dir_path = ''

    filePath = mc.fileDialog2(fm=1,
                              dir=dir_path,
                              ff='Maya Files (*.ma *.mb);;Maya ASCII (*.ma);;Maya Binary (*.mb);;All Files (*.*)')

    if not filePath:
        return ''

    if fullPath:
        return filePath[0]
    else:
        return os.path.basename(filePath[0])



def referenceFile(filePath, namespace):
    if not os.path.isfile(filePath):
        return False

    if filePath in mc.file(q=True, r=True):
        return False

    mc.file(filePath, r=True, ns=namespace)
    return True



def referenceRig(filePath):
    referenceFile(filePath, 'Rig')


def referenceMocap(filePath):
    referenceFile(filePath, 'Mocap')


def referenceHumanIK(filePath):
    referenceFile(filePath, 'HumanIK')



def transferMocapData(rig_path, mocap_path, human_path):
    '''
    '''
    #- create new scene
    if not mel.eval('saveChanges("file -new -f")'):
        return

    #- close polygon objects display...
    currentPannel = mc.getPanel(wf=True)
    mc.modelEditor(currentPannel, e=True, alo=False)
    mc.modelEditor(currentPannel, e=True, j=True)

    #- reference files
    referenceRig(rig_path)
    referenceMocap(mocap_path)
    referenceHumanIK(human_path)

    mc.currentUnit(t='pal')
    #- get mocap objects...
    mocap_locator  = mc.listRelatives(mc.ls(mc.referenceQuery(mocap_path, n=True), typ='locator'), p=True, path=True)
    mocap_joints   = ' '.join(mc.ls(mc.referenceQuery(mocap_path, n=True), typ='joint'))
    mocap_HIK_node = mc.ls(mc.referenceQuery(mocap_path, n=True), typ='HIKCharacterNode')


    #- initialize Human IK...
    mel.eval('HIKCharacterControlsTool;')


    #- create mocap humanIK...
    if mocap_HIK_node:
        mocap_HIK_node = mocap_HIK_node[0]

    else:
        mocap_hip = ''
        humanIK_hip = ''

        mocap_nodes = mc.referenceQuery(mocap_path, n=True)
        for jnt in mocap_nodes:
            if mc.nodeType(jnt) != 'joint':
                continue
            mc.setAttr(jnt + '.rx', 0)
            mc.setAttr(jnt + '.ry', 0)
            mc.setAttr(jnt + '.rz', 0)

            if re.search('hips$', jnt, re.I):
                mocap_hip = jnt

        humanIK_hip = mocap_hip.replace('Mocap:', 'HumanIK:')

        posi = mc.xform(humanIK_hip, q=True, ws=True, t=True)
        mc.xform(mocap_hip, ws=True, t=posi)

        old_humanIK_nodes = mc.ls(typ='HIKCharacterNode')
        mel.eval('hikCreateDefinition;')
        mocap_HIK_node = [x for x in mc.ls(typ='HIKCharacterNode') if x not in old_humanIK_nodes][0]
        mocap_HIK_node = mc.rename(mocap_HIK_node, 'mocap_HIK')

        #- defind mocap hunmanIk skeleton...
        if mocap_locator:
            mel.eval('setCharacterObject("{0}", "{1}", {2}, 0);'.format(mocap_locator[0], mocap_HIK_node, 0))

        #- reading data
        f = open(os.path.join(scriptTool.getScriptPath(), 'match_data_skeleton.json'), 'r')
        skeleton_data = json.load(f)
        f.close()

        for name, index in skeleton_data.iteritems():
            joint_name = re.search('\S*{0}\S*'.format(name), mocap_joints)
            if not joint_name:
                continue
            mel.eval('setCharacterObject("{0}", "{1}", {2}, 0);'.format(joint_name.group(), mocap_HIK_node, index))

        #- lock defintion...
        mel.eval('HIKCharacterControlsTool;hikUpdateDefinitionUI();')
        if not mel.eval('hikCheckDefinitionLocked("{0}")'.format(mocap_HIK_node)):
            mel.eval('hikToggleLockDefinition();')


    #-
    HumanIK_HIK_node = mc.ls(mc.referenceQuery(human_path, n=True), typ='HIKCharacterNode')
    if not HumanIK_HIK_node:
        raise RuntimeError('Can\'t find HumanIK in file - {0} '.format(human_path))


    #- set connect...
    mc.optionMenuGrp('hikCharacterList', e=True, v=HumanIK_HIK_node[0])
    mel.eval('hikToggleLockDefinition();hikUpdateCurrentCharacterFromUI(); hikUpdateContextualUI();')

    mc.optionMenuGrp('hikSourceList', e=True, v=' ' + mocap_HIK_node)
    mel.eval("hikUpdateCurrentSourceFromUI(); hikUpdateContextualUI();")


    #- bake keys...
    mocap_keys = mc.keyframe(mc.referenceQuery(mocap_path, dp=True, n=True), q=True)
    start_time = min(mocap_keys) - 1
    end_time   = max(mocap_keys) + 1


    mc.playbackOptions(e=True, min=start_time, max=end_time, ast=start_time, aet=end_time)
    mel.eval('hikBakeToControlRig 0;')


    #- go to Tpose...
    humanIK_eft_nodes = mc.ls(mc.referenceQuery(human_path, dp=True, n=True), typ='hikIKEffector')
    mc.select(humanIK_eft_nodes)
    mel.eval('hikStancePose;')


    #- constraint rig controls...
    f = open(os.path.join(scriptTool.getScriptPath(), 'match_data_control.json'), 'r')
    control_data = json.load(f)
    f.close()

    humanIK_file_nodes = ' '.join(mc.ls(mc.referenceQuery(human_path, dp=True, n=True), typ=('joint', 'hikIKEffector')))
    rig_controls = list()
    pole_vector_control = dict()
    for k, v in control_data.iteritems():
        src_control = re.search('{0}:\w+{1}'.format(mc.file(human_path, q=True, ns=True), k), humanIK_file_nodes)
        dst_control = mc.ls('{0}:*{1}'.format(mc.file(rig_path, q=True, ns=True), v))

        if not src_control:
            continue

        src_control = src_control.group()
        if not mc.objExists(src_control):
            continue

        if not dst_control:
            continue

        dst_control = dst_control[0]
        if not mc.objExists(dst_control):
            continue

        try:
            mc.pointConstraint(src_control, dst_control, mo=True)
        except RuntimeError:
            pass

        try:
            mc.orientConstraint(src_control, dst_control, mo=True)
        except RuntimeError:
            pass

        if re.search('ElbowIKC|KneeIKC|pole', dst_control):
            pole_vector_control[src_control] = dst_control

        rig_controls.append(dst_control)

    #- bake keys
    mc.bakeResults(rig_controls, sm=True, t=(start_time, end_time), sb=1, dic=True, pok=True, sac=False, ral=False, bol=False, mr=True, cp=False, s=False)

    #- delete constraint
    mc.delete(mc.listRelatives(rig_controls, c=True, typ=('pointConstraint', 'parentConstraint')))

    #- fix pole vector...
    for i in range(int(start_time), int(end_time) + 1):
        mc.currentTime(i)
        for src, dst in pole_vector_control.iteritems():
            posi = mc.xform(src, q=True, ws=True, t=True)
            base_posi = mc.xform(dst, q=True, ws=True, rp=True)
            mx, my, mz = posi[0]-base_posi[0], posi[1]-base_posi[1], posi[2]-base_posi[2]

            mc.move(mx, my, mz, dst, r=True)
            mc.setKeyframe(dst)

    #- show polygon curves and nurbssurfaces...
    mc.modelEditor(currentPannel, e=True, nc=True, ns=True, pm=True)

    #-
    return True